package com.pccc.batch.config.remotepartition.handler;

import javax.sql.DataSource;

import org.springframework.batch.core.Step;
import org.springframework.batch.core.StepExecutionListener;
import org.springframework.batch.core.explore.JobExplorer;
import org.springframework.batch.core.partition.StepExecutionSplitter;
import org.springframework.batch.core.partition.support.Partitioner;
import org.springframework.batch.core.partition.support.StepExecutionAggregator;
import org.springframework.batch.core.repository.JobRepository;
import org.springframework.batch.core.step.builder.PartitionStepBuilder;
import org.springframework.batch.core.step.builder.StepBuilder;
import org.springframework.batch.integration.partition.MessageChannelPartitionHandler;
import org.springframework.batch.integration.partition.RemotePartitioningManagerStepBuilder;
import org.springframework.beans.factory.BeanCreationException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.integration.aggregator.HeaderAttributeCorrelationStrategy;
import org.springframework.integration.aggregator.SimpleSequenceSizeReleaseStrategy;
import org.springframework.integration.channel.QueueChannel;
import org.springframework.integration.core.MessagingTemplate;
import org.springframework.integration.dsl.IntegrationFlows;
import org.springframework.integration.dsl.StandardIntegrationFlow;
import org.springframework.integration.dsl.context.IntegrationFlowContext;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.PollableChannel;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.util.Assert;

public class PollerRemotePartitioningManagerStepBuilder extends PartitionStepBuilder {

	
	private static final long DEFAULT_POLL_INTERVAL = 10000L;
	private static final long DEFAULT_TIMEOUT = -1L;

	private MessagingTemplate messagingTemplate;
	private MessageChannel inputChannel;
	private MessageChannel outputChannel;
	private JobExplorer jobExplorer;
	private BeanFactory beanFactory;
	private long pollInterval = DEFAULT_POLL_INTERVAL;
	private long timeout = DEFAULT_TIMEOUT;
	private DataSource dataSource;
	
	public PollerRemotePartitioningManagerStepBuilder(String stepName) {
		super(new StepBuilder(stepName));
	}

	
	public Step build() {
		Assert.state(this.outputChannel == null || this.messagingTemplate == null,
				"You must specify either an outputChannel or a messagingTemplate but not both.");

		// configure messaging template
		if (this.messagingTemplate == null) {
			this.messagingTemplate = new MessagingTemplate();
			this.messagingTemplate.setDefaultChannel(this.outputChannel);
			if (this.logger.isDebugEnabled()) {
				this.logger.debug("No messagingTemplate was provided, using a default one");
			}
		}

		// Configure the partition handler
		final MessageChannelPartitionHandler partitionHandler = new MessageChannelPartitionHandler();
		partitionHandler.setStepName(getStepName());
		partitionHandler.setGridSize(getGridSize());
		partitionHandler.setMessagingOperations(this.messagingTemplate);

		if (isPolling()) {
			partitionHandler.setJobExplorer(this.jobExplorer);
			partitionHandler.setPollInterval(this.pollInterval);
			partitionHandler.setTimeout(this.timeout);
			partitionHandler.setDataSource(this.dataSource);
		}
		else {
			PollableChannel replies = new QueueChannel();
			partitionHandler.setReplyChannel(replies);
			StandardIntegrationFlow standardIntegrationFlow = IntegrationFlows
					.from(this.inputChannel)
					.aggregate(aggregatorSpec -> aggregatorSpec.releaseStrategy(new SimpleSequenceSizeReleaseStrategy())
							.correlationStrategy(new HeaderAttributeCorrelationStrategy("correlationid")))
					.channel(replies)
					.get();
			IntegrationFlowContext integrationFlowContext = this.beanFactory.getBean(IntegrationFlowContext.class);
			integrationFlowContext.registration(standardIntegrationFlow)
					.autoStartup(false)
					.register();
		}

		try {
			partitionHandler.afterPropertiesSet();
			super.partitionHandler(partitionHandler);
		}
		catch (Exception e) {
			throw new BeanCreationException("Unable to create a manager step for remote partitioning", e);
		}

		return super.build();
	}
	
	private boolean isPolling() {
		return this.inputChannel == null;
	}
	
	
	public PollerRemotePartitioningManagerStepBuilder dataSource(DataSource dataSource) {
		this.dataSource = dataSource;
		return this;
	}
	
	/**
	 * Set the input channel on which replies from workers will be received.
	 * @param inputChannel the input channel
	 * @return this builder instance for fluent chaining
	 */
	public PollerRemotePartitioningManagerStepBuilder inputChannel(MessageChannel inputChannel) {
		Assert.notNull(inputChannel, "inputChannel must not be null");
		this.inputChannel = inputChannel;
		return this;
	}

	/**
	 * Set the output channel on which requests to workers will be sent. By using
	 * this setter, a default messaging template will be created and the output
	 * channel will be set as its default channel.
	 * <p>Use either this setter or {@link RemotePartitioningManagerStepBuilder#messagingTemplate(MessagingTemplate)}
	 * to provide a fully configured messaging template.</p>
	 *
	 * @param outputChannel the output channel.
	 * @return this builder instance for fluent chaining
	 * @see RemotePartitioningManagerStepBuilder#messagingTemplate(MessagingTemplate)
	 */
	public PollerRemotePartitioningManagerStepBuilder outputChannel(MessageChannel outputChannel) {
		Assert.notNull(outputChannel, "outputChannel must not be null");
		this.outputChannel = outputChannel;
		return this;
	}

	/**
	 * Set the {@link MessagingTemplate} to use to send data to workers.
	 * <strong>The default channel of the messaging template must be set</strong>.
	 * <p>Use either this setter to provide a fully configured messaging template or
	 * provide an output channel through {@link RemotePartitioningManagerStepBuilder#outputChannel(MessageChannel)}
	 * and a default messaging template will be created.</p>
	 *
	 * @param messagingTemplate the messaging template to use
	 * @return this builder instance for fluent chaining
	 * @see RemotePartitioningManagerStepBuilder#outputChannel(MessageChannel)
	 */
	public PollerRemotePartitioningManagerStepBuilder messagingTemplate(MessagingTemplate messagingTemplate) {
		Assert.notNull(messagingTemplate, "messagingTemplate must not be null");
		this.messagingTemplate = messagingTemplate;
		return this;
	}

	/**
	 * Set the job explorer.
	 * @param jobExplorer the job explorer to use.
	 * @return this builder instance for fluent chaining
	 */
	public PollerRemotePartitioningManagerStepBuilder jobExplorer(JobExplorer jobExplorer) {
		Assert.notNull(jobExplorer, "jobExplorer must not be null");
		this.jobExplorer = jobExplorer;
		return this;
	}

	/**
	 * How often to poll the job repository for the status of the workers. Defaults to 10 seconds.
	 * @param pollInterval the poll interval value in milliseconds
	 * @return this builder instance for fluent chaining
	 */
	public PollerRemotePartitioningManagerStepBuilder pollInterval(long pollInterval) {
		Assert.isTrue(pollInterval > 0, "The poll interval must be greater than zero");
		this.pollInterval = pollInterval;
		return this;
	}

	/**
	 * When using job repository polling, the time limit to wait. Defaults to -1 (no timeout).
	 * @param timeout the timeout value in milliseconds
	 * @return this builder instance for fluent chaining
	 */
	public PollerRemotePartitioningManagerStepBuilder timeout(long timeout) {
		this.timeout = timeout;
		return this;
	}

	/**
	 * Set the bean factory.
	 * @param beanFactory the bean factory to use
	 * @return this builder instance for fluent chaining
	 */
	public PollerRemotePartitioningManagerStepBuilder beanFactory(BeanFactory beanFactory) {
		this.beanFactory = beanFactory;
		return this;
	}
	
	
	@Override
	public PollerRemotePartitioningManagerStepBuilder repository(JobRepository jobRepository) {
		super.repository(jobRepository);
		return this;
	}

	@Override
	public PollerRemotePartitioningManagerStepBuilder transactionManager(PlatformTransactionManager transactionManager) {
		super.transactionManager(transactionManager);
		return this;
	}

	@Override
	public PollerRemotePartitioningManagerStepBuilder partitioner(String workerStepName, Partitioner partitioner) {
		super.partitioner(workerStepName, partitioner);
		return this;
	}

	@Override
	public PollerRemotePartitioningManagerStepBuilder gridSize(int gridSize) {
		super.gridSize(gridSize);
		return this;
	}

	@Override
	public PollerRemotePartitioningManagerStepBuilder step(Step step) {
		super.step(step);
		return this;
	}

	@Override
	public PollerRemotePartitioningManagerStepBuilder splitter(StepExecutionSplitter splitter) {
		super.splitter(splitter);
		return this;
	}

	@Override
	public PollerRemotePartitioningManagerStepBuilder aggregator(StepExecutionAggregator aggregator) {
		super.aggregator(aggregator);
		return this;
	}

	@Override
	public PollerRemotePartitioningManagerStepBuilder startLimit(int startLimit) {
		super.startLimit(startLimit);
		return this;
	}

	@Override
	public PollerRemotePartitioningManagerStepBuilder listener(Object listener) {
		super.listener(listener);
		return this;
	}

	@Override
	public PollerRemotePartitioningManagerStepBuilder listener(StepExecutionListener listener) {
		super.listener(listener);
		return this;
	}

	@Override
	public PollerRemotePartitioningManagerStepBuilder allowStartIfComplete(boolean allowStartIfComplete) {
		super.allowStartIfComplete(allowStartIfComplete);
		return this;
	}
	
}
